import uuid
import json
import yaml
from datetime import datetime
from typing import Self

from pydantic import BaseModel, Field, model_validator


class MemoryNote(BaseModel):
    id: str = Field(default_factory=lambda: str(uuid.uuid4()))
    # documents
    keywords: list[str] = Field(default_factory=list)
    summary: str = Field(default="General")
    tags: list[str] = Field(default_factory=list)
    content: str = Field(default="")
    # metadatas
    links: list[str] = Field(default_factory=list)
    retrieval_count: int = Field(default=0)
    timestamp: str = Field(default="")
    last_accessed: str = Field(default="")

    @model_validator(mode="after")
    def get_current_time(self) -> Self:
        if not self.timestamp:
            current_time = datetime.now().strftime("%Y%m%d%H%M%S")
            self.timestamp = current_time
            self.last_accessed = current_time
        else:
            self.last_accessed = self.timestamp
        return self


def serialize_to_document(document_dict : dict | MemoryNote) -> str:
    if isinstance(document_dict, MemoryNote):
        return f"""content: {document_dict.content}
keywords: {document_dict.keywords}
summary: {document_dict.summary}
tags: {document_dict.tags}"""
    else:
        return f"""content: {document_dict["content"]}
keywords: {document_dict["keywords"]}
summary: {document_dict["summary"]}
tags: {document_dict["tags"]}"""

def deserialize_from_document(document: str) -> dict:
    return yaml.safe_load(document)

def serialize(note: MemoryNote) -> tuple[str, str, dict]:
    document = f"""content: {note.content}
keywords: {note.keywords}
summary: {note.summary}
tags: {note.tags}"""
    metadata = dict(
        links=json.dumps(note.links),
        retrieval_count=note.retrieval_count,
        timestamp=note.timestamp,
        last_accessed=note.last_accessed
    )
    return note.id, document, metadata


def deserialize(id: str, document: str, metadata: dict) -> MemoryNote:
    document_dict = yaml.safe_load(document)
    metadata["links"] = json.loads(metadata["links"])
    return MemoryNote(id=id, **document_dict, **metadata)
